import os
import json
import logging
import networkx as nx
from collections import Counter
from datetime import datetime

def setup_logging():
    """设置日志配置"""
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    log_dir = os.path.join("output", "logs")
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    
    log_file = os.path.join(log_dir, f'behavior_graph_{timestamp}.log')
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s [%(levelname)s] %(message)s',
        handlers=[
            logging.FileHandler(log_file),
            logging.StreamHandler()
        ]
    )

def json_to_behavior_graph(json_file):
    """将Cuckoo沙箱的JSON文件转换为更丰富的行为图"""
    try:
        with open(json_file, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        G = nx.DiGraph()
        processes = {}
        
        # 添加基本信息节点
        info = data.get('info', {})
        main_node = f"main_{info.get('id', 'unknown')}"
        G.add_node(main_node, 
                   type='main',
                   category=info.get('category', ''),
                   score=info.get('score', 0),
                   platform=info.get('platform', ''))
        
        # 处理签名信息
        for sig in data.get('signatures', []):
            sig_node = f"signature_{sig['name']}"
            G.add_node(sig_node, 
                      type='signature',
                      name=sig['name'],
                      severity=sig.get('severity', 0),
                      description=sig.get('description', ''))
            G.add_edge(main_node, sig_node, type='has_signature')
            
            # 处理签名中的标记
            for mark in sig.get('marks', []):
                if mark['type'] == 'call':
                    call = mark['call']
                    call_node = f"call_{mark['cid']}"
                    G.add_node(call_node,
                             type='api_call',
                             api=call['api'],
                             category=call['category'],
                             status=call['status'])
                    G.add_edge(sig_node, call_node, type='contains_call')
                    
                    # 添加进程关联
                    proc_node = f"process_{mark['pid']}"
                    if not G.has_node(proc_node):
                        G.add_node(proc_node, type='process', pid=mark['pid'])
                    G.add_edge(proc_node, call_node, type='makes_call')
                
                elif mark['type'] == 'ioc':
                    ioc_node = f"ioc_{hash(mark['ioc'])}"
                    G.add_node(ioc_node,
                             type='ioc',
                             category=mark['category'],
                             value=mark['ioc'])
                    G.add_edge(sig_node, ioc_node, type='identifies')
        
        # 处理进程行为
        if 'behavior' in data and 'processes' in data['behavior']:
            for proc in data['behavior']['processes']:
                pid = str(proc.get('process_id', ''))
                ppid = str(proc.get('parent_id', ''))
                proc_name = proc.get('process_name', 'unknown')
                
                proc_node = f"process_{pid}"
                G.add_node(proc_node, 
                          type='process',
                          name=proc_name,
                          pid=pid)
                
                if ppid and ppid != pid:
                    parent_node = f"process_{ppid}"
                    if not G.has_node(parent_node):
                        G.add_node(parent_node, type='process', pid=ppid)
                    G.add_edge(parent_node, proc_node, type='parent')
                
                # 处理API调用
                for call in proc.get('calls', []):
                    call_node = f"call_{call.get('id', hash(str(call)))}"
                    G.add_node(call_node,
                             type='api_call',
                             api=call['api'],
                             category=call['category'])
                    G.add_edge(proc_node, call_node, type='makes_call')
                    
                    # 处理调用参数
                    args = call.get('arguments', {})
                    for arg_name, arg_value in args.items():
                        if isinstance(arg_value, (str, int)) and str(arg_value).strip():
                            arg_node = f"arg_{hash(str(arg_value))}"
                            G.add_node(arg_node,
                                     type='argument',
                                     name=arg_name,
                                     value=str(arg_value))
                            G.add_edge(call_node, arg_node, type='has_argument')
        
        # 检查图的大小
        if G.number_of_nodes() == 0 or G.number_of_edges() == 0:
            logging.warning(f"生成的图为空 (文件: {json_file})")
            return None
            
        logging.info(f"成功生成图 (文件: {json_file})")
        logging.info(f"- 节点数: {G.number_of_nodes()}")
        logging.info(f"- 边数: {G.number_of_edges()}")
        logging.info(f"- 节点类型统计: {dict(Counter(nx.get_node_attributes(G, 'type').values()))}")
        
        return G
    
    except Exception as e:
        logging.error(f"处理JSON文件 {json_file} 时出错: {str(e)}")
        return None

def process_json_files(json_dir, output_dir):
    """处理目录中的所有JSON文件并生成图"""
    setup_logging()
    
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    processed = 0
    empty = 0
    errors = 0
    
    for file in os.listdir(json_dir):
        if not file.endswith('.json'):
            continue
            
        json_path = os.path.join(json_dir, file)
        graph_path = os.path.join(output_dir, file.replace('.json', '.graphml'))
        
        try:
            G = json_to_behavior_graph(json_path)
            if G is not None:
                nx.write_graphml(G, graph_path)
                processed += 1
            else:
                empty += 1
        except Exception as e:
            logging.error(f"处理文件 {file} 时出错: {str(e)}")
            errors += 1
    
    logging.info(f"处理完成:")
    logging.info(f"- 成功处理: {processed}")
    logging.info(f"- 空图: {empty}")
    logging.info(f"- 错误: {errors}")
    
    return processed, empty, errors 